/*******************************************************\
|* Author: Trevor Sundberg
\*******************************************************/

// Defines
#define WIN32_LEAN_AND_MEAN

// Includes
#include "SkugoConnection.hpp"
#include "SkugoDebug.hpp"
#include <WinSock2.h>
#include <concurrent_queue.h>
#include <vector>

// Defines
#define SkugoInfiniteLoop for(;;)

// Include the winsock libraries
#pragma comment(lib, "Ws2_32.lib")

namespace Skugo
{
  // Using directives
  using std::vector;

  // The Win32 implementation of local connection (IConnection)
  class Win32LocalConnection : public IConnection
  {
  public:

    //***************************************************************************
    // A basic chunk that is written and read
    struct Chunk
    {
      char*    Data;
      sizeType Size;
    };

    //***************************************************************************
    // Type-defines
    typedef Concurrency::concurrent_queue<Chunk> ChunkQueue;

    //***************************************************************************
    // Thread numbers
    enum Threads
    {
      SendThreadNum = 0,
      ReceiveThreadNum = 1,
      NumThreads
    };

    //***************************************************************************
    // Constructor
    Win32LocalConnection(SOCKET socket) :
      mSocket(socket)
    {
      // Create the send thread
      mThreadHandle[SendThreadNum] = CreateThread(
        NULL,
        0,
        SendThread,
        this,
        0,
        &mThreadId[SendThreadNum]);

      // Create the receive thread
      mThreadHandle[ReceiveThreadNum] = CreateThread(
        NULL,
        0,
        ReceiveThread,
        this,
        0,
        &mThreadId[ReceiveThreadNum]);
    }

    //***************************************************************************
    // Destructor
    ~Win32LocalConnection()
    {
      // Disconnect the connection
      this->Disconnect();
    }

    //***************************************************************************
    // A thread function that is responsible for sending data
    static DWORD WINAPI SendThread(LPVOID parameter)
    {
      // Get the 'self' object from the thread parameter
      Win32LocalConnection* self = (Win32LocalConnection*)parameter;

      // Send data until the socket is closed (or has an error)
      SkugoInfiniteLoop
      {
        // A packet that we'll read data into
        Chunk outgoing;

        // Pop data off the outgoing queue
        if (self->mOutgoing.try_pop(outgoing))
        {
          // Start off at the front of the packet
          sizeType currentPosition = 0;

          // Loop until we reach the end of the packet...
          while (currentPosition < outgoing.Size)
          {
            // Send the data!
            int bytesSent = send(self->mSocket, outgoing.Data + currentPosition, outgoing.Size, 0);

            // If we hit a socket error...
            SkugoReturnIf(bytesSent == SOCKET_ERROR, 0, "A socket sending error occurred");

            // Move forward by the number of bytes we sent
            currentPosition += bytesSent;
          }

          // Delete the packet data
          delete[] outgoing.Data;
        }
      }
    }

    //***************************************************************************
    // A thread function that is responsible for receiving data
    static DWORD WINAPI ReceiveThread(LPVOID parameter)
    {
      // Get the 'self' object from the thread parameter
      Win32LocalConnection* self = (Win32LocalConnection*)parameter;

      // Receive data until the socket is closed (or has an error)
      SkugoInfiniteLoop
      {
        // Declare a buffer that will hold the read in messages
        const sizeType BufferSize = 8192;
        char buffer[BufferSize];

        // Read data from the socket
        int bytesRead = recv(self->mSocket, buffer, BufferSize, 0);

        // If we had an error we need to exit out...
        SkugoReturnIf(bytesRead == SOCKET_ERROR, 0, "A socket receiving error occurred");

        // If the connection closed gracefully then also exit out
        if (bytesRead == 0)
          return 0;

        // Create a new packet
        Chunk incoming;

        // Copy the data into a new packet buffer
        incoming.Data = new char[bytesRead];
        memcpy(incoming.Data, buffer, bytesRead);
        
        // Set the buffer size to the number of bytes we read
        incoming.Size = bytesRead;

        // Add the incoming packet to the incoming queue
        self->mIncoming.push(incoming);
      }
    }

    //***************************************************************************
    virtual void Write(const char* dataIn, sizeType size)
    {
      // Create a new packet
      Chunk outgoing;

      // Copy the data into a new buffer
      outgoing.Data = new char[size];
      memcpy(outgoing.Data, dataIn, size);

      // Set the size of the outgoing packet
      outgoing.Size = size;

      // Add the buffer to the outgoing queue
      mOutgoing.push(outgoing);
    }

    //***************************************************************************
    virtual bool Read(char** dataOut, sizeType* sizeOut)
    {
      // Create a new packet
      Chunk incoming;

      // Add the buffer to the outgoing queue
      if (mIncoming.try_pop(incoming))
      {
        // If we got data, output it and return success
        *dataOut = incoming.Data;
        *sizeOut = incoming.Size;
        return true;
      }

      // Otherwise, return that no data has arrived
      return false;
    }

    //***************************************************************************
    virtual void Disconnect()
    {
      // If the socket is already invalid, don't bother doing anything
      if (mSocket == INVALID_SOCKET)
        return;

      // Shutdown the socket and close it
      shutdown(mSocket, SD_SEND);
      closesocket(mSocket);

      // Wait for the send and receive threads to converge
      WaitForMultipleObjects(NumThreads, mThreadHandle, TRUE, INFINITE);

      // Clear the threads and ids
      mThreadHandle[SendThreadNum]    = INVALID_HANDLE_VALUE;
      mThreadHandle[ReceiveThreadNum] = INVALID_HANDLE_VALUE;
      mThreadId[SendThreadNum]        = 0;
      mThreadId[ReceiveThreadNum]     = 0;

      // Clear the socket
      mSocket = INVALID_SOCKET;

      // While the outgoing queue still has data in it...
      while(mOutgoing.empty() == false)
      {
        // Get the next entry in the outgoing data
        Chunk outgoing;
        if (mOutgoing.try_pop(outgoing))
        {
          // Delete the data
          delete[] outgoing.Data;
        }
      }

      // While the incoming queue still has data in it...
      while(mIncoming.empty() == false)
      {
        // Get the next entry in the outgoing data
        Chunk incoming;
        if (mIncoming.try_pop(incoming))
        {
          // Delete the data
          delete[] incoming.Data;
        }
      }
    }

    //***************************************************************************
    virtual bool IsConnected()
    {
      if (mSocket != INVALID_SOCKET)
      {
        // HACK Perform a 'select' here
        return true;
      }

      return false;
    }

    //***************************************************************************
    virtual bool IsLocal()
    {
      SkugoWarning("We need an actual valid check to see if the connection is local!");
      return false;
    }

  private:

    // Store the send and receive threads (in that order)
    DWORD mThreadId[NumThreads];
    HANDLE mThreadHandle[NumThreads];

    // Store the connected socket
    SOCKET mSocket;

    // All the outgoing data
    ChunkQueue mOutgoing;

    // All the incoming data
    ChunkQueue mIncoming;
  };

  // The Win32 implementation of IHost
  class Win32Host : public IHost
  {
  public:

    // Type-defines
    typedef Concurrency::concurrent_queue<Win32LocalConnection*> ConnectionQueue;
    typedef vector<Win32LocalConnection*> ConnectionList;

    // Constructor
    Win32Host() :
      mSocket(INVALID_SOCKET)
    {
      // Startup winsock
      WSADATA data;
      int result = WSAStartup(MAKEWORD(2, 2), &data);

      // An error just in case it fails
      SkugoErrorIf(result != 0, "An error occurred while trying to startup WinSock");
    }

    // Destructor
    ~Win32Host()
    {
      // Disconnect the socket (only if it's connected)
      StopHosting();

      // Close all the connections
      CloseAllConnections();

      // Cleanup winsock
      WSACleanup();
    }

    // A thread function that accepts the incoming connection and starst the send/receive
    static DWORD WINAPI SocketThread(LPVOID parameter)
    {
      // Get the 'self' object from the thread parameter
      Win32Host* self = (Win32Host*)parameter;

      // Store socket address information
      sockaddr_in sockAddress;
      memset(&sockAddress, 0, sizeof(sockAddress));
      int sockAddrSize = sizeof(sockaddr_in);

      // Accept sockets until the socket is closed...
      SkugoInfiniteLoop
      {
        // Accept any incoming connections
        SOCKET newSocket = accept(self->mSocket, (SOCKADDR*) &sockAddress, &sockAddrSize);

        // If we got an invalid socket, break out
        if (newSocket == INVALID_SOCKET)
          break;

        // Add a new connection to the latest connections
        self->mNewConnections.push(new Win32LocalConnection(newSocket));
      }

      // Return a success
      return 0;
    }

    // Start hosting a session
    virtual void StartHosting(unsigned short port)
    {
      // First, we should disconnect (it won't do anything if the socket isn't valid)
      StopHosting();

      // Create a sock-addr struct and fill it in with the details of how we bind
      sockaddr_in service;
      service.sin_family      = AF_INET;
      service.sin_addr.s_addr = INADDR_ANY;
      service.sin_port        = htons(port);

      // Create the socket
      mSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);

      // Bind the socket to the given port
      bind(mSocket, (SOCKADDR*) &service, sizeof(service));

      // Tell the socket to listen for a connection
      listen(mSocket, SOMAXCONN);

      // Create a thread that we'll use to do all the socket operations
      mThreadHandle = CreateThread(
        NULL,
        0,
        SocketThread,
        this,
        0,
        &mThreadId);
    }

    // Stop the hosting (so no more clients can connect)
    virtual void StopHosting()
    {
      // If the socket is already invalid, don't bother doing anything
      if (mSocket == INVALID_SOCKET)
        return;

      // Shutdown the socket and close it
      shutdown(mSocket, SD_SEND);
      closesocket(mSocket);

      // Wait for the thread to terminate
      WaitForSingleObject(mThreadHandle, INFINITE);

      // Clear the thread and id
      mThreadHandle = INVALID_HANDLE_VALUE;
      mThreadId = 0;

      // Clear the socket
      mSocket = INVALID_SOCKET;
    }

    // Gets the next recent connection (or returns NULL)
    virtual IConnection* PopNewConnection()
    {
      // What we'll store the connection in
      Win32LocalConnection* newConnection = NULL;

      // If we are able to pop...
      if (mNewConnections.try_pop(newConnection))
      {
        // Add the connection to the list of active connections
        mActiveConnections.push_back(newConnection);

        // Return the connection that we popped!
        return newConnection;
      }

      // Otherwise there was nothing
      return NULL;
    }

    // Write data to the client
    virtual void Write(const char* dataIn, sizeType sizeIn)
    {
      // Loop through all the active connections
      for (sizeType i = 0; i < mActiveConnections.size(); ++i)
      {
        // Write data out to the socket
        mActiveConnections[i]->Write(dataIn, sizeIn);
      }
    }

    // Close all the connections
    virtual void CloseAllConnections()
    {
      // Loop through all the active connections
      while (mActiveConnections.empty() == false)
      {
		// Delete the connection (which also disconnects it)
		delete mActiveConnections.back();

		// Remove the connection
		mActiveConnections.pop_back();
      }

	  // Loop through all pending connections
	  while (IConnection* connection = PopNewConnection())
	  {
		// Delete the connection (which also disconnects it)
		delete connection;
	  }
    }

    // Stops hosting entirely, and closes all active or pending connections
    virtual void Shutdown()
	{
		// Stop hosting and close the connections
		StopHosting();
		CloseAllConnections();
	}

  private:

    // The thread id
    DWORD mThreadId;

    // The thread handle
    HANDLE mThreadHandle;

    // Store the socket that we use to accept connections
    SOCKET mSocket;

    // All the recent connections (these connections are coming from another thread)
    ConnectionQueue mNewConnections;

    // All the active connections (these connects are on the main thread)
    ConnectionList mActiveConnections;
  };

  // Create a default implementation of the IHost interface
  IHost* IHost::Create()
  {
    return new Win32Host();
  }
}
